import os
import json
import hashlib
import time
import base64
import logging
from logging.handlers import RotatingFileHandler
from flask import Flask, request, jsonify, Response, stream_template
import requests
from dotenv import load_dotenv
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import struct

# 加载环境变量
load_dotenv('config.env')

# 配置日志
def setup_logging():
    """配置日志系统"""
    # 创建logs目录
    log_dir = 'logs'
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    
    # 配置日志格式
    log_format = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建logger
    logger = logging.getLogger('wechat_robot')
    log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
    logger.setLevel(getattr(logging, log_level))
    
    # 清除现有的handlers
    logger.handlers.clear()
    
    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(log_format)
    logger.addHandler(console_handler)
    
    # 文件处理器 - 应用日志
    max_size = int(os.getenv('LOG_MAX_SIZE', 10*1024*1024))  # 默认10MB
    backup_count = int(os.getenv('LOG_BACKUP_COUNT', 5))
    
    app_handler = RotatingFileHandler(
        os.path.join(log_dir, 'app.log'),
        maxBytes=max_size,
        backupCount=backup_count,
        encoding='utf-8'
    )
    app_handler.setLevel(logging.INFO)
    app_handler.setFormatter(log_format)
    logger.addHandler(app_handler)
    
    # 文件处理器 - 错误日志
    error_handler = RotatingFileHandler(
        os.path.join(log_dir, 'error.log'),
        maxBytes=max_size,
        backupCount=backup_count,
        encoding='utf-8'
    )
    error_handler.setLevel(logging.ERROR)
    error_handler.setFormatter(log_format)
    logger.addHandler(error_handler)
    
    return logger

# 初始化日志
logger = setup_logging()

app = Flask(__name__)

class WeChatCrypto:
    """企业微信消息加解密类"""
    
    def __init__(self, token, encoding_aes_key, corp_id):
        self.token = token
        self.encoding_aes_key = encoding_aes_key
        self.corp_id = corp_id
        self.aes_key = base64.b64decode(encoding_aes_key + '=')
    
    def check_signature(self, signature, timestamp, nonce, msg_encrypt):
        tmp_arr = [self.token, timestamp, nonce, msg_encrypt]
        tmp_arr.sort()
        tmp_str = ''.join(tmp_arr)
        hash_str = hashlib.sha1(tmp_str.encode('utf-8')).hexdigest()
        is_valid = hash_str == signature
        return is_valid
    
    def decrypt(self, text):
        try:
            encrypted_data = base64.b64decode(text)
            cipher = AES.new(self.aes_key, AES.MODE_CBC, self.aes_key[:16])
            decrypted = cipher.decrypt(encrypted_data)
            pad_len = decrypted[-1]
            content = decrypted[:-pad_len]
            random_str = content[:16]
            msg_len = struct.unpack("!I", content[16:20])[0]
            json_bytes = content[20:20+msg_len]
            corp_id_bytes = content[20+msg_len:]
            expected_corp_id = corp_id_bytes.decode('utf-8')
            if expected_corp_id != self.corp_id:
                logger.warning(f"企业ID不匹配: 期望={expected_corp_id}, 实际={self.corp_id}")
            result = json_bytes.decode('utf-8')
            return result
        except Exception as e:
            logger.error(f"消息解密失败: {e}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return None
    
    def encrypt(self, reply_msg, timestamp, nonce):
        try:
            json_bytes = reply_msg.encode('utf-8')
            msg_len = struct.pack("!I", len(json_bytes))
            random_str = os.urandom(16)
            text = random_str + msg_len + json_bytes + self.corp_id.encode('utf-8')
            text = pad(text, AES.block_size)
            cipher = AES.new(self.aes_key, AES.MODE_CBC, self.aes_key[:16])
            cipher_text = cipher.encrypt(text)
            base64_text = base64.b64encode(cipher_text).decode('utf-8')
            tmp_arr = [self.token, timestamp, nonce, base64_text]
            tmp_arr.sort()
            tmp_str = ''.join(tmp_arr)
            signature = hashlib.sha1(tmp_str.encode('utf-8')).hexdigest()
            return base64_text, signature
        except Exception as e:
            logger.error(f"消息加密失败: {e}")
            return None, None

class DifyAPI:
    """Dify API 接口类"""
    
    def __init__(self, api_url, api_key):
        self.api_url = api_url
        self.api_key = api_key
        self.headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
    
    def send_message_stream(self, message, conversation_id=None, user_id=None):
        """发送消息到Dify并返回流式响应"""
        try:
            # 根据Dify API文档构建请求体，使用流式模式
            payload = {
                "inputs": {
                    "accessToken": self.api_key  # 添加accessToken参数
                },
                "query": message,
                "response_mode": "streaming",  # 使用流式模式
                "user": user_id or "default_user"
            }
            
            if conversation_id:
                payload["conversation_id"] = conversation_id
            
            # 添加连接超时和读取超时
            response = requests.post(
                f"{self.api_url}/chat-messages",
                headers=self.headers,
                json=payload,
                timeout=(10, 60),  # (连接超时, 读取超时)
                stream=True  # 启用流式传输
            )
            
            if response.status_code == 200:
                return response
            else:
                logger.error(f"Dify API 流式请求失败: {response.status_code} - {response.text}")
                return None
                
        except requests.exceptions.ConnectionError as e:
            logger.error(f"Dify API连接失败: {e}")
            logger.error("请检查Dify API URL是否正确，以及网络连接是否正常")
            return None
        except requests.exceptions.Timeout as e:
            logger.error(f"Dify API请求超时: {e}")
            return None
        except Exception as e:
            logger.error(f"调用Dify API流式接口时发生错误: {e}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return None

    def send_message(self, message, conversation_id=None, user_id=None):
        """发送消息到Dify（兼容模式）"""
        try:
            # 根据Dify API文档构建请求体
            payload = {
                "inputs": {
                    "accessToken": self.api_key  # 添加accessToken参数
                },
                "query": message,
                "response_mode": "blocking",
                "user": user_id or "default_user"
            }
            
            if conversation_id:
                payload["conversation_id"] = conversation_id
            
            response = requests.post(
                f"{self.api_url}/chat-messages",
                headers=self.headers,
                json=payload,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                answer = result.get('answer', '抱歉，我暂时无法回答您的问题。')
                return answer
            else:
                logger.error(f"Dify API 请求失败: {response.status_code} - {response.text}")
                return "抱歉，服务暂时不可用，请稍后再试。"
                
        except Exception as e:
            logger.error(f"调用Dify API时发生错误: {e}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return "抱歉，服务暂时不可用，请稍后再试。"

# 初始化配置
wechat_crypto = WeChatCrypto(
    token=os.getenv('WECHAT_TOKEN'),
    encoding_aes_key=os.getenv('WECHAT_ENCODING_AES_KEY'),
    corp_id=os.getenv('WECHAT_CORP_ID')
)

dify_api = DifyAPI(
    api_url=os.getenv('DIFY_API_URL'),
    api_key=os.getenv('DIFY_API_KEY')
)

# 内存缓存，用于存储流式消息内容
stream_cache = {}

@app.route('/wechat', methods=['GET', 'POST'])
def wechat_webhook():
    """企业微信webhook接口"""
    try:
        # 获取请求参数
        msg_signature = request.args.get('msg_signature', '')
        timestamp = request.args.get('timestamp', '')
        nonce = request.args.get('nonce', '')
        
        # 添加详细的请求日志
        logger.info(f"收到企业微信请求:")
        logger.info(f"  方法: {request.method}")
        logger.info(f"  URL参数: msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}")
        logger.info(f"  请求头: {dict(request.headers)}")
        
        if request.method == 'GET':
            # 处理验证请求
            echostr = request.args.get('echostr', '')
            logger.info(f"GET请求 - echostr: {echostr}")
            
            # 验证签名
            if not wechat_crypto.check_signature(msg_signature, timestamp, nonce, echostr):
                logger.error("企业微信验证失败 - 签名不匹配")
                return '签名验证失败', 403
            
            # 如果echostr是加密的，需要解密
            try:
                decrypted_echostr = wechat_crypto.decrypt(echostr)
                if decrypted_echostr:
                    return decrypted_echostr
            except:
                # 如果解密失败，返回原始echostr
                return echostr
            return echostr
        
        elif request.method == 'POST':
            # 处理POST消息
            logger.info(f"POST请求 - Content-Type: {request.content_type}")
            
            # 记录原始请求数据
            raw_data = request.get_data()
            logger.info(f"原始请求数据: {raw_data}")
            
            # 尝试解析JSON
            try:
                json_data = request.get_json()
                logger.info(f"解析的JSON数据: {json_data}")
            except Exception as e:
                logger.error(f"JSON解析失败: {e}")
                # 尝试解析为字符串
                try:
                    text_data = raw_data.decode('utf-8')
                    logger.info(f"请求体文本: {text_data}")
                except:
                    logger.error("无法解析请求体")
                return '无效的JSON数据', 400
            
            if not json_data:
                logger.error("无法解析JSON数据")
                return '无效的JSON数据', 400
            
            # 检查所有可能的加密字段名
            encrypt_msg = None
            possible_fields = ['Encrypt', 'encrypt', 'MsgEncrypt', 'msg_encrypt']
            
            for field in possible_fields:
                if field in json_data:
                    encrypt_msg = json_data[field]
                    logger.info(f"找到加密字段 '{field}': {encrypt_msg[:50]}..." if encrypt_msg else "None")
                    break
            
            # 如果没有找到加密字段，检查是否是明文消息
            if not encrypt_msg:
                logger.info(f"未找到加密字段，检查是否为明文消息。可用字段: {list(json_data.keys())}")
                
                # 检查是否是明文消息格式
                if 'msgtype' in json_data and 'from' in json_data:
                    logger.info("检测到明文消息格式，直接处理")
                    message_data = json_data
                else:
                    logger.error(f"未找到加密消息内容，可用字段: {list(json_data.keys())}")
                    return '缺少加密消息内容', 400
            else:
                # 验证签名 - 使用Encrypt字段进行签名验证
                if not wechat_crypto.check_signature(msg_signature, timestamp, nonce, encrypt_msg):
                    logger.error("企业微信POST消息验证失败 - 签名不匹配")
                    return '签名验证失败', 403
                
                # 解密消息
                decrypted_json = wechat_crypto.decrypt(encrypt_msg)
                if not decrypted_json:
                    logger.error("消息解密失败")
                    return '消息解密失败', 400
                
                logger.info(f"解密后的消息: {decrypted_json}")
                
                # 解析消息
                try:
                    message_data = json.loads(decrypted_json)
                except json.JSONDecodeError as e:
                    logger.error(f"JSON解析失败: {e}")
                    return 'JSON解析失败', 400
            
            # 提取消息信息 - 使用正确的字段名
            msg_type = message_data.get('msgtype', '')
            from_user = message_data.get('from', {}).get('userid', '')
            msg_id = message_data.get('msgid', '')
            user_message = message_data.get('text', {}).get('content', '')
            
            logger.info(f"解析消息成功:")
            logger.info(f"  消息类型: {msg_type}")
            logger.info(f"  发送者: {from_user}")
            logger.info(f"  消息ID: {msg_id}")
            logger.info(f"  消息内容: {user_message}")
            
            # 检查是否是流式消息刷新请求 - 根据官方文档修正
            if msg_type == 'stream':
                logger.info(f"收到流式消息刷新请求")
                
                # 从消息中提取流式消息ID - 根据官方文档修正
                stream_id = message_data.get('stream', {}).get('id', '')
                if not stream_id:
                    logger.error("流式消息刷新请求缺少stream.id")
                    return '缺少stream.id', 400
                
                # 从缓存中获取流式消息内容
                cached_content = stream_cache.get(stream_id)
                if not cached_content:
                    logger.error(f"未找到流式消息缓存: {stream_id}")
                    return '流式消息不存在', 404
                
                # 构建流式消息响应
                stream_message = {
                    'msgtype': 'stream',
                    'stream': {
                        'id': stream_id,
                        'finish': True,
                        'content': cached_content
                    }
                }
                
                # 根据消息类型决定是否加密响应
                if encrypt_msg:
                    # 加密响应
                    reply_json = json.dumps(stream_message, ensure_ascii=False)
                    encrypted_reply, signature = wechat_crypto.encrypt(reply_json, timestamp, nonce)
                    
                    if not encrypted_reply:
                        logger.error("流式消息响应加密失败")
                        return '响应加密失败', 500
                    
                    # 构建响应
                    response_data = {
                        'encrypt': encrypted_reply,
                        'msg_signature': signature,
                        'timestamp': timestamp,
                        'nonce': nonce
                    }
                else:
                    # 明文响应
                    response_data = stream_message
                
                logger.info("返回流式消息刷新响应")
                return jsonify(response_data)
            
            # 处理文本消息
            elif msg_type == 'text':
                logger.info(f"处理文本消息 - 内容: {user_message}")
                
                # 生成流式消息ID
                stream_id = f"stream_{timestamp}_{from_user}"
                
                # 调用Dify API获取流式响应
                dify_response = dify_api.send_message_stream(user_message, user_id=from_user)
                
                if not dify_response:
                    # 如果流式请求失败，使用普通请求
                    dify_answer = dify_api.send_message(user_message, user_id=from_user)
                    
                    # 构建普通文本消息响应
                    text_message = {
                        'msgtype': 'text',
                        'text': {
                            'content': dify_answer
                        }
                    }
                    
                    # 根据消息类型决定是否加密响应
                    if encrypt_msg:
                        # 加密响应
                        reply_json = json.dumps(text_message, ensure_ascii=False)
                        encrypted_reply, signature = wechat_crypto.encrypt(reply_json, timestamp, nonce)
                        
                        if not encrypted_reply:
                            logger.error("文本消息响应加密失败")
                            return '响应加密失败', 500
                        
                        # 构建响应
                        response_data = {
                            'encrypt': encrypted_reply,
                            'msg_signature': signature,
                            'timestamp': timestamp,
                            'nonce': nonce
                        }
                    else:
                        # 明文响应
                        response_data = text_message
                    
                    return jsonify(response_data)
                
                # 处理流式响应
                def generate_stream():
                    try:
                        full_content = ""
                        line_count = 0
                        
                        for line in dify_response.iter_lines(decode_unicode=True):
                            line_count += 1
                            
                            if not line:
                                continue
                            
                            if line.startswith('data: '):
                                data_str = line[6:]  # 移除 'data: ' 前缀
                                
                                if data_str == '[DONE]':
                                    break
                                
                                try:
                                    data = json.loads(data_str)
                                except json.JSONDecodeError:
                                    continue
                                
                                # 提取answer字段
                                answer = None
                                if data.get('event') == 'message':
                                    answer = data.get('data', {}).get('answer', '')
                                elif data.get('event') == 'node_finished':
                                    answer = data.get('data', {}).get('outputs', {}).get('answer', '')
                                elif data.get('event') == 'workflow_finished':
                                    answer = data.get('data', {}).get('outputs', {}).get('answer', '')
                                
                                if answer and answer not in full_content:
                                    full_content += answer
                                    
                                    # 构建流式消息
                                    stream_message = {
                                        'msgtype': 'stream',
                                        'stream': {
                                            'id': stream_id,
                                            'finish': False,
                                            'content': answer
                                        }
                                    }
                                    
                                    # 根据消息类型决定是否加密响应
                                    if encrypt_msg:
                                        # 加密响应
                                        reply_json = json.dumps(stream_message, ensure_ascii=False)
                                        encrypted_reply, signature = wechat_crypto.encrypt(reply_json, timestamp, nonce)
                                        
                                        if encrypted_reply:
                                            response_data = {
                                                'encrypt': encrypted_reply,
                                                'msg_signature': signature,
                                                'timestamp': timestamp,
                                                'nonce': nonce
                                            }
                                        else:
                                            continue
                                    else:
                                        # 明文响应
                                        response_data = stream_message
                                    
                                    yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"
                        
                        # 缓存完整内容
                        if full_content:
                            stream_cache[stream_id] = full_content
                        
                        # 发送最终完成消息
                        final_stream_message = {
                            'msgtype': 'stream',
                            'stream': {
                                'id': stream_id,
                                'finish': True,
                                'content': full_content
                            }
                        }
                        
                        # 根据消息类型决定是否加密响应
                        if encrypt_msg:
                            reply_json = json.dumps(final_stream_message, ensure_ascii=False)
                            encrypted_reply, signature = wechat_crypto.encrypt(reply_json, timestamp, nonce)
                            
                            if encrypted_reply:
                                response_data = {
                                    'encrypt': encrypted_reply,
                                    'msg_signature': signature,
                                    'timestamp': timestamp,
                                    'nonce': nonce
                                }
                            else:
                                return
                        else:
                            # 明文响应
                            response_data = final_stream_message
                        
                        yield f"data: {json.dumps(response_data, ensure_ascii=False)}\n\n"
                    
                    except Exception as e:
                        logger.error(f"流式响应处理失败: {e}")
                        import traceback
                        logger.error(f"详细错误信息: {traceback.format_exc()}")
                
                # 返回流式响应
                return Response(
                    generate_stream(),
                    mimetype='text/plain',
                    headers={
                        'Cache-Control': 'no-cache',
                        'Connection': 'keep-alive',
                        'Content-Type': 'text/plain; charset=utf-8'
                    }
                )
            
            else:
                # 忽略其他类型的消息
                logger.info(f"忽略非文本消息 - 类型: {msg_type}")
                return 'success'
    
    except Exception as e:
        logger.error(f"处理企业微信消息时发生错误: {e}")
        import traceback
        logger.error(f"详细错误信息: {traceback.format_exc()}")
        return '服务器内部错误', 500

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({
        'status': 'healthy',
        'timestamp': time.time(),
        'service': 'wechat_robot'
    })

@app.route('/', methods=['GET'])
def index():
    """根路径"""
    return jsonify({
        'service': 'WeChat Enterprise Robot',
        'version': '1.0.0',
        'status': 'running'
    })

if __name__ == '__main__':
    port = int(os.getenv('PORT', 53230))
    debug = os.getenv('FLASK_DEBUG', 'False').lower() == 'true'
    
    app.run(
        host='0.0.0.0',
        port=port,
        debug=debug,
        threaded=True
    ) 